
import os

xla_flags = os.environ.get("XLA_FLAGS", "")
xla_flags += " --xla_gpu_triton_gemm_any=True"
os.environ["XLA_FLAGS"] = xla_flags
os.environ["MUJOCO_GL"] = "egl"
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"

import functools
import jax
import jax.numpy as jp
import mediapy as media
from brax.training.agents.ppo import networks as ppo_networks
from brax.training.agents.ppo import train as ppo
from etils import epath

from mujoco_playground import BraxEnvWrapper, locomotion

# Enable persistent compilation cache.
jax.config.update("jax_compilation_cache_dir", "/tmp/jax_cache")
jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)
jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)

env_name = "Go1JoystickFlatTerrain"
env_cfg = locomotion.get_default_config(env_name)
env = locomotion.load(env_name, config=env_cfg)
jit_reset = jax.jit(env.reset)
jit_step = jax.jit(env.step)

CKPT_PATH = "/home/kevin/dev/mjx_control/learning/checkpoints/Go1JoystickFlatTerrain-20241112-170242"
CKPT_PATH = epath.Path(CKPT_PATH)
assert CKPT_PATH.exists()

latest_ckpts = list(CKPT_PATH.glob("*"))
if not latest_ckpts:
  raise ValueError("No checkpoints found")
# Remove .json files.
latest_ckpts = [
    ckpt for ckpt in latest_ckpts if not ckpt.name.endswith(".json")
]
latest_ckpts.sort(key=lambda x: int(x.name))
print(f"Found {len(latest_ckpts)} checkpoints")
latest_ckpt = latest_ckpts[-1]
restore_checkpoint_path = latest_ckpt
print(f"Restoring from {restore_checkpoint_path}")

make_networks_factory = functools.partial(
    ppo_networks.make_ppo_networks,
    policy_hidden_layer_sizes=(128, 128, 128, 128),
)

train_fn = functools.partial(
    ppo.train,
    num_timesteps=0,
    episode_length=1000,
    normalize_observations=True,
    restore_checkpoint_path=restore_checkpoint_path,
    network_factory=make_networks_factory,
    num_envs=2,
)
make_inference_fn, params, _ = train_fn(environment=BraxEnvWrapper(env))

inference_fn = make_inference_fn(params, deterministic=True)
jit_inference_fn = jax.jit(inference_fn)

eval_env = locomotion.load(env_name, config=env_cfg)
jit_reset = jax.jit(eval_env.reset)
jit_step = jax.jit(eval_env.step)
#%%
cmd = jp.array([0.3, 0, 0.1])
rng = jax.random.PRNGKey(12345)
rollout = []
for _ in range(1):
  rng, reset_rng = jax.random.split(rng)
  state = jit_reset(reset_rng)
  state.info["command"] = cmd
  for i in range(env_cfg.episode_length):
    act_rng, rng = jax.random.split(rng)
    ctrl, _ = jit_inference_fn(state.obs, act_rng)
    state = jit_step(state, ctrl)
    state.info["command"] = cmd
    rollout.append(state)
#%%
render_every = 2
fps = 1.0 / eval_env.dt / render_every
traj = rollout[::render_every]

frames = eval_env.render(
    traj,
    camera="track",
    height=480,
    width=640,
)
media.show_video(frames, fps=fps, loop=False)
#%%
